import os
import logging
import boto3


from label_studio_ml.model import LabelStudioMLBase
from label_studio_ml.utils import get_image_size, get_single_tag_keys
from label_studio.core.settings.base import DATA_UNDEFINED_NAME
from label_studio.core.utils.io import json_load, get_data_dir
from botocore.exceptions import ClientError
from urllib.parse import urlparse
from label_studio_ml.yolov4_tiny.test import run_test

logger = logging.getLogger(__name__)


class YoloV3Detection(LabelStudioMLBase):
    """Object detector based on https://github.com/open-mmlab/mmdetection"""

    def __init__(self, image_dir=None, labels_file=None, **kwargs):
        """
        Load MMDetection model from config and checkpoint into memory.
        (Check https://mmdetection.readthedocs.io/en/v1.2.0/GETTING_STARTED.html#high-level-apis-for-testing-images)

        Optionally set mappings from COCO classes to target labels
        :param config_file: Absolute path to MMDetection config file (e.g. /home/user/mmdetection/configs/faster_rcnn/faster_rcnn_r50_fpn_1x.py)
        :param checkpoint_file: Absolute path MMDetection checkpoint file (e.g. /home/user/mmdetection/checkpoints/faster_rcnn_r50_fpn_1x_20181010-3d1b3351.pth)
        :param image_dir: Directory where images are stored (should be used only in case you use direct file upload into Label Studio instead of URLs)
        :param labels_file: file with mappings from COCO labels to custom labels {"airplane": "Boeing"}
        :param score_threshold: score threshold to wipe out noisy results
        :param device: device (cpu, cuda:0, cuda:1, ...)
        :param kwargs:
        """
        super(YoloV3Detection, self).__init__(**kwargs)

        self.root_path = os.path.join(os.path.realpath(__file__), '../')
        self.checkpoint_file = '/var/run/label-studio-ml-backend/label_studio_ml/examples/yolov4tiny_detection/yolov4_tiny_final.ckpt'
        self.labels_file = labels_file
        # default Label Studio image upload folder
        upload_dir = os.path.join(get_data_dir(), 'media', 'upload')
        self.image_dir = "/outputs/label-studio/media/upload"
        logger.debug(f'{self.__class__.__name__} reads images from {self.image_dir}')
        if self.labels_file and os.path.exists(self.labels_file):
            self.label_map = json_load(self.labels_file)
        else:
            self.label_map = {}

        self.from_name, self.to_name, self.value, self.labels_in_config = get_single_tag_keys(
            self.parsed_label_config, 'RectangleLabels', 'Image')
        #self.labels_in_config = set(self.labels_in_config)

    def _get_image_url(self, task):
        image_url = task['data'].get(self.value) or task['data'].get(DATA_UNDEFINED_NAME)
        if image_url.startswith('s3://'):
            # presign s3 url
            r = urlparse(image_url, allow_fragments=False)
            bucket_name = r.netloc
            key = r.path.lstrip('/')
            client = boto3.client('s3')
            try:
                image_url = client.generate_presigned_url(
                    ClientMethod='get_object',
                    Params={'Bucket': bucket_name, 'Key': key}
                )
            except ClientError as exc:
                logger.warning(f'Can\'t generate presigned URL for {image_url}. Reason: {exc}')
        return image_url

    def predict(self, tasks, **kwargs):
        assert len(tasks) == 1
        task = tasks[0]
        image_url = self._get_image_url(task)
        project_id = image_url.split('/')[-2]

        image_path = os.path.join(self.image_dir, project_id, os.path.basename(image_url))
        print(f"image_path: {image_path}")

        predicts = run_test(self.checkpoint_file, image_path)
        results = []
        all_scores = []
        img_width, img_height = get_image_size(image_path)
        for p in predicts:
            print(p)
            output_label = self.labels_in_config[p['category_id'] - 1]

            x, y, xmax, ymax = p['bbox']
            results.append({
                'from_name': self.from_name,
                'to_name': self.to_name,
                'type': 'rectanglelabels',
                'value': {
                    'rectanglelabels': [output_label],
                    'x': float(x / img_width * 100),
                    'y': float(y / img_height * 100),
                    'width': float(xmax / img_width * 100),
                    'height': float(ymax / img_height * 100)
                },
                'score': float(p['score'])
            })
            all_scores.append(float(p['score']))

        avg_score = sum(all_scores) / max(len(all_scores), 1)
        return [{
            'result': results,
            'score': avg_score
        }]
